// 实验最终版：阻抗模型 + 状态机（距离控制）+ TD软启动
// 描述：
// 1. Z轴采用阻抗控制维持恒定接触力。
// 2. 接触瞬间，使用TD（跟踪微分器）对目标力进行平滑，避免冲击。
// 3. 接触判定基于力/力变化率，并带有滞回，更加鲁棒。
// 4. 稳定接触后，延时2秒，然后沿X轴按指定距离往复运动。
// 5. 所有关键参数（阻抗、TD、运动速度/距离）均可通过ROS参数在线动态调整。

#include <rclcpp/rclcpp.hpp>

// MoveIt Servo
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <tf2_ros/buffer.h>
#include <rcl_interfaces/msg/set_parameters_result.hpp>

#include <atomic>
#include <algorithm>
#include <cmath>
#include <string>
#include <chrono>

using namespace std::chrono_literals;

namespace imp_node {  // 避免与其他源文件的全局符号冲突

// ========== 状态机定义 ==========
enum class ControlState {
  CONTACT_AND_MAINTAIN_FORCE, // 初始状态：接触并维持恒力
  DELAY_BEFORE_MOVE,        // 状态2：达到力之后，延时
  MOVE_POSITIVE_X,          // 状态3：沿X轴正向运动
  MOVE_NEGATIVE_X           // 状态4：沿X轴负向运动
};

// ========== 小工具 ==========
static inline double sgn(double x) { return (x > 0) - (x < 0); }
static inline double clamp_scalar(double v, double lo, double hi) {
  return v < lo ? lo : (v > hi ? hi : v);
}
static inline std::string state_to_string(ControlState state) {
    switch (state) {
        case ControlState::CONTACT_AND_MAINTAIN_FORCE: return "CONTACT_FORCE";
        case ControlState::DELAY_BEFORE_MOVE:        return "DELAY";
        case ControlState::MOVE_POSITIVE_X:          return "MOVE_POS_X";
        case ControlState::MOVE_NEGATIVE_X:          return "MOVE_NEG_X";
        default:                                     return "UNKNOWN";
    }
}


// ========== 节点句柄/通信 ==========
static rclcpp::Node::SharedPtr node_;
static rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_pub_;
static rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;
static rclcpp::node_interfaces::OnSetParametersCallbackHandle::SharedPtr param_cb_handle;

static const rclcpp::Logger LOGGER = rclcpp::get_logger("plane_impedance");

// ========== 参数（可在线调） ==========
// 目标恒力与增益
static std::atomic<double> g_f_target{5.0};
static std::atomic<double> g_kappa{1.0};

// 阻抗参数
static std::atomic<double> g_md{1.0};
static std::atomic<double> g_bd{100.0};
static std::atomic<double> g_kd{0.0};

// 限幅
static std::atomic<double> g_max_vel{0.4};
static std::atomic<double> g_max_acc{1.5};

// 横向运动参数
static std::atomic<double> g_lateral_vel{0.05};
static std::atomic<double> g_lateral_distance{0.3};

// [新增] TD/CONTACT 参数
static std::atomic<double> g_f_enter{0.5};    // 进入接触力阈值
static std::atomic<double> g_f_exit{0.3};     // 退出接触力阈值
static std::atomic<double> g_df_enter{15.0};  // 进入接触力斜率阈值 (N/s)
static std::atomic<bool>   g_td_enable{true}; // 是否启用 TD
static std::atomic<double> g_td_r0{8.0};      // TD “最速”参数 r0

// Twist frame
static std::string g_twist_frame{"base_link"};

// ========== 状态变量 ==========
static std::atomic<double> g_force_meas{0.0};
static double x_ = 0.0;
static double xdot_ = 0.0;
static rclcpp::Time t_prev_;
static bool has_prev_time_ = false;

// 状态机变量
static std::atomic<ControlState> g_current_state{ControlState::CONTACT_AND_MAINTAIN_FORCE};
static rclcpp::Time g_state_transition_time;
static double g_lateral_displacement_ = 0.0;

// [新增] TD/CONTACT 状态变量
static bool   in_contact_ = false;
static bool   has_prevF_  = false;
static double Fe_prev_    = 0.0;
static rclcpp::Time tF_prev_;
static double g_td_v1 = 0.0;
static double g_td_v2 = 0.0;
static bool   g_td_inited = false;

// ======================
// [新增] TD核心函数
// ======================
static inline double fhan(double x1_minus_ref, double x2, double r, double h)
{
  const double d  = r * h * h;
  const double a0 = h * x2;
  const double y  = x1_minus_ref + a0;
  const double a1 = std::sqrt(d * (d + 8.0 * std::fabs(y)));
  const double a2 = a0 + sgn(y) * (a1 - d) * 0.5;
  return -r * ( (std::fabs(a2) > d) ? sgn(a2) : (a2 / d) );
}
static inline void td_update(double Ref, double r0, double h, double &v1, double &v2)
{
  const double f0 = fhan(v1 - Ref, v2, r0, h);
  v1 += h * v2;
  v2 += h * f0;
}

// ========== 订阅 /sim_force ==========
// [修改] forceCallback 仅用于更新力，状态切换移至 controlTimer
static void forceCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
{
  g_force_meas.store(msg->wrench.force.z, std::memory_order_relaxed);
}

// ========== 定时器回调：主控制逻辑 ==========
static void controlTimer()
{
  const rclcpp::Time now = node_->now();
  double dt = 0.01;
  if (has_prev_time_) {
    dt = std::max(1e-4, (now - t_prev_).seconds());
  }
  t_prev_ = now;
  has_prev_time_ = true;

  const double Fe   = g_force_meas.load(std::memory_order_relaxed);
  const double Fref = g_f_target.load();

  // =================================================================
  // 1. [新增] 接触判定与TD软启动
  // =================================================================
  double dFdt = 0.0;
  if (has_prevF_) {
    double dts = std::max(1e-4, (now - tF_prev_).seconds());
    dFdt = (Fe - Fe_prev_) / dts;
  }
  Fe_prev_ = Fe; tF_prev_ = now; has_prevF_ = true;

  bool last_contact_state = in_contact_;
  if (!in_contact_) {
    if (std::fabs(Fe) >= g_f_enter.load() || std::fabs(dFdt) >= g_df_enter.load()) {
      in_contact_ = true;
      g_td_v1 = Fe; g_td_v2 = 0.0; g_td_inited = true; // 初始化TD：从当前实际力开始平滑
    }
  } else {
    if (std::fabs(Fe) <= g_f_exit.load() && std::fabs(dFdt) <= 0.5 * g_df_enter.load()) {
      in_contact_ = false;
      g_td_inited = false;
    }
  }

  // [修改] 使用新的接触判定来触发状态机
  if (in_contact_ && !last_contact_state && g_current_state.load() == ControlState::CONTACT_AND_MAINTAIN_FORCE) {
      RCLCPP_INFO(LOGGER, "已接触! 从 [CONTACT_FORCE] 切换到 [DELAY]。");
      g_state_transition_time = now;
      g_current_state.store(ControlState::DELAY_BEFORE_MOVE);
  }

  double Fref_eff = Fref; // 默认使用原始目标力
  if (in_contact_ && g_td_enable.load()) {
    if (!g_td_inited) { // 双保险
      g_td_v1 = Fe; g_td_v2 = 0.0; g_td_inited = true;
    }
    td_update(Fref, g_td_r0.load(), dt, g_td_v1, g_td_v2);
    Fref_eff = g_td_v1; // 使用TD平滑后的目标力
  }

  // =================================================================
  // 2. Z轴阻抗控制（使用平滑后的 Fref_eff）
  // =================================================================
  const double eF = Fref_eff - Fe;
  const double u  = g_kappa.load() * eF;

  const double md = std::max(1e-6, g_md.load());
  const double bd = g_bd.load();
  const double kd = g_kd.load();
  double xdd = (u - bd * xdot_ - kd * x_) / md;
  xdd = clamp_scalar(xdd, -g_max_acc.load(), g_max_acc.load());

  double xdot_new = xdot_ + xdd * dt;
  xdot_new = clamp_scalar(xdot_new, -g_max_vel.load(), g_max_vel.load());
  xdot_ = xdot_new;
  x_   += xdot_ * dt;

  // =================================================================
  // 3. X轴往复运动状态机
  // =================================================================
  double lateral_vx = 0.0;
  ControlState current_state = g_current_state.load();
  const double target_dist = g_lateral_distance.load();

  switch (current_state) {
    case ControlState::CONTACT_AND_MAINTAIN_FORCE:
      break;
    case ControlState::DELAY_BEFORE_MOVE:
      if ((now - g_state_transition_time).seconds() >= 2.0) {
        RCLCPP_INFO(LOGGER, "延时结束，开始X轴运动。目标距离: %.3f m", target_dist);
        g_lateral_displacement_ = 0.0;
        g_current_state.store(ControlState::MOVE_POSITIVE_X);
      }
      break;
    case ControlState::MOVE_POSITIVE_X:
      lateral_vx = g_lateral_vel.load();
      g_lateral_displacement_ += lateral_vx * dt;
      if (g_lateral_displacement_ >= target_dist) {
        g_lateral_displacement_ = 0.0;
        g_current_state.store(ControlState::MOVE_NEGATIVE_X);
      }
      break;
    case ControlState::MOVE_NEGATIVE_X:
      lateral_vx = -g_lateral_vel.load();
      g_lateral_displacement_ += std::abs(lateral_vx * dt);
      if (g_lateral_displacement_ >= target_dist) {
        g_lateral_displacement_ = 0.0;
        g_current_state.store(ControlState::MOVE_POSITIVE_X);
      }
      break;
  }

  // =================================================================
  // 4. 发布最终的Twist指令
  // =================================================================
  auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
  msg->header.stamp = now;
  msg->header.frame_id = g_twist_frame;
  msg->twist.linear.x = lateral_vx;
  msg->twist.linear.y = 0.0;
  msg->twist.linear.z = -xdot_;
  twist_pub_->publish(std::move(msg));

  // 节流打印
  RCLCPP_INFO_THROTTLE(
      LOGGER, *node_->get_clock(), 50,
      "MODE:%s | STATE:%-12s | Fe=%.2f Fref_eff=%.2f | z_dot=%.3f x_dot=%.3f | x_dist=%.3f/%.3f",
      in_contact_ ? "CONTACT" : "FREE",
      state_to_string(current_state).c_str(), Fe, Fref_eff, -xdot_, lateral_vx, g_lateral_displacement_, target_dist);
}

} // namespace imp_node

int main(int argc, char** argv)
{
  using namespace imp_node;

  rclcpp::init(argc, argv);
  rclcpp::NodeOptions opts;
  opts.use_intra_process_comms(false);
  node_ = std::make_shared<rclcpp::Node>("plane_impedance_td", opts);

  // —— MoveIt/Servo 设置 (与之前相同) ——
  auto tf_buffer = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
  auto psm = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
      node_, "robot_description", tf_buffer, "planning_scene_monitor");
  if (psm->getPlanningScene()) {
    psm->startStateMonitor("/joint_states");
    psm->setPlanningScenePublishingFrequency(25);
    psm->startPublishingPlanningScene(
        planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE, "/moveit_servo/publish_planning_scene");
    psm->startSceneMonitor();
    psm->providePlanningSceneService();
  } else {
    RCLCPP_ERROR(LOGGER, "Planning scene 未正确配置");
    return EXIT_FAILURE;
  }
  auto servo_params = moveit_servo::ServoParameters::makeServoParameters(node_);
  auto servo = std::make_unique<moveit_servo::Servo>(node_, servo_params, psm);
  servo->start();

  // —— Pub/Sub ——
  twist_pub_ = node_->create_publisher<geometry_msgs::msg::TwistStamped>("~/delta_twist_cmds", 10);
  force_sub_ = node_->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/sim_force", rclcpp::SensorDataQoS(), &forceCallback);

  // —— 参数声明 ——
  node_->declare_parameter<double>("f_target", g_f_target.load());
  node_->declare_parameter<double>("kappa", g_kappa.load());
  node_->declare_parameter<double>("imp.md", g_md.load());
  node_->declare_parameter<double>("imp.bd", g_bd.load());
  node_->declare_parameter<double>("imp.kd", g_kd.load());
  node_->declare_parameter<double>("limit.max_vel", g_max_vel.load());
  node_->declare_parameter<double>("limit.max_acc", g_max_acc.load());
  node_->declare_parameter<std::string>("twist_frame", g_twist_frame);
  node_->declare_parameter<double>("lateral_velocity", g_lateral_vel.load());
  node_->declare_parameter<double>("lateral_distance", g_lateral_distance.load());
  // [新增] TD/CONTACT 参数声明
  node_->declare_parameter<double>("contact.f_enter", g_f_enter.load());
  node_->declare_parameter<double>("contact.f_exit",  g_f_exit.load());
  node_->declare_parameter<double>("contact.df_enter", g_df_enter.load());
  node_->declare_parameter<bool>("td.enable",   g_td_enable.load());
  node_->declare_parameter<double>("td.r0",     g_td_r0.load());

  // —— 同步初值 ——
  g_f_target.store(node_->get_parameter("f_target").as_double());
  g_kappa.store(node_->get_parameter("kappa").as_double());
  g_md.store(node_->get_parameter("imp.md").as_double());
  g_bd.store(node_->get_parameter("imp.bd").as_double());
  g_kd.store(node_->get_parameter("imp.kd").as_double());
  g_max_vel.store(node_->get_parameter("limit.max_vel").as_double());
  g_max_acc.store(node_->get_parameter("limit.max_acc").as_double());
  g_twist_frame = node_->get_parameter("twist_frame").as_string();
  g_lateral_vel.store(node_->get_parameter("lateral_velocity").as_double());
  g_lateral_distance.store(node_->get_parameter("lateral_distance").as_double());
  // [新增] TD/CONTACT 参数同步
  g_f_enter.store(node_->get_parameter("contact.f_enter").as_double());
  g_f_exit.store(node_->get_parameter("contact.f_exit").as_double());
  g_df_enter.store(node_->get_parameter("contact.df_enter").as_double());
  g_td_enable.store(node_->get_parameter("td.enable").as_bool());
  g_td_r0.store(node_->get_parameter("td.r0").as_double());

  // —— 动态参数回调 ——
  param_cb_handle = node_->add_on_set_parameters_callback(
    [](const std::vector<rclcpp::Parameter>& params) {
      rcl_interfaces::msg::SetParametersResult res; res.successful = true;
      for (const auto& p : params) {
        const auto &name = p.get_name();
        if      (name == "f_target")         { imp_node::g_f_target.store(p.as_double()); }
        else if (name == "kappa")            { imp_node::g_kappa.store(p.as_double()); }
        else if (name == "imp.md")           { imp_node::g_md.store(p.as_double()); }
        else if (name == "imp.bd")           { imp_node::g_bd.store(p.as_double()); }
        else if (name == "imp.kd")           { imp_node::g_kd.store(p.as_double()); }
        else if (name == "lateral_velocity") { imp_node::g_lateral_vel.store(p.as_double()); }
        else if (name == "lateral_distance") { imp_node::g_lateral_distance.store(p.as_double()); }
        // [新增] TD/CONTACT 动态更新
        else if (name == "contact.f_enter")  { imp_node::g_f_enter.store(p.as_double()); }
        else if (name == "contact.f_exit")   { imp_node::g_f_exit.store(p.as_double()); }
        else if (name == "contact.df_enter") { imp_node::g_df_enter.store(p.as_double()); }
        else if (name == "td.enable")        { imp_node::g_td_enable.store(p.as_bool()); }
        else if (name == "td.r0")            { imp_node::g_td_r0.store(p.as_double()); }
      }
      return res;
    }
  );

  // —— 控制循环：10ms（100Hz） ——
  auto timer = node_->create_wall_timer(10ms, &controlTimer);
  auto exec = std::make_unique<rclcpp::executors::MultiThreadedExecutor>();
  exec->add_node(node_);
  exec->spin();

  rclcpp::shutdown();
  return 0;
}
